import os
import random
import numpy


class label_info:
    def __init__(self, loc_index, name):
        self.loc_index = loc_index
        self.name = name
        self.Files = []

def enum_files(path,Files):
        for _, _, files in os.walk(path):
            for file in files:
                if file.endswith(".npy"):
                    Files.append(file)



class audioDatastore:

    def __init__(self,labels):

        self.DataLoc = None
        self.Labels = labels
        self.Data = [[] for i in range(len(labels))]

    def add_dataLoc(self, path ):
        self.DataLoc=path

        for root, dirs, files in os.walk(path):
            for dir in dirs:
                if dir in self.Labels:
                    n = self.Labels.index(dir)

                    _path = os.path.join(root, dir)
                    enum_files(_path,self.Data[n])
                else:
                    print("%s not listed in classes.ignored."%dir)


    def get_full_path(self, i, j):

        label = self.Labels[i]
        fname = self.Data[i][j]
        return os.path.join(self.DataLoc, label, fname)

    def gen_dataset(self):
        trainloader = []
        for i in range(len(self.Labels)):
            for j in range(len(self.Data[i])):
                trainloader.append((i, j))
        return trainloader

    def load(self, i, j):
        fn = self.get_full_path(i, j)
        data = None

        try:
            data = numpy.load(fn)
        except IOError:
            print('Error While Opening the file!%s' % fn)

        if numpy.any(numpy.logical_or(numpy.isnan(data),numpy.isinf(data))):
            print("NaN.(%s)"%fn)
            data = None
        return data


def load_classes(path):
    cls_fn = os.path.join(path,"classes.txt")
    classes=[]
    if os.path.exists(cls_fn):
        with open(cls_fn,"r") as f:
            for line in f:
                cls=line.strip()
                if cls not in classes:
                    classes.append(cls)
    return classes

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        sys.exit(0)

    ads = audioDatastore()
    ads.add_dataLoc(sys.argv[1])

    for i in range(len(ads.Labels)):
        for j in range(min(5, len(ads.Labels[i].Files))):
            s = ads.get_full_path(i, j)
            print(s)

    trainloader, testloader = ads.loader()
    print(len(trainloader))
    print(len(testloader))
    for k in range(min(5, len(trainloader))):
        i, j = trainloader[k]
        s = ads.get_full_path(i, j)
        print(s)
    print()
    for k in range(min(5, len(testloader))):
        i, j = testloader[k]
        s = ads.get_full_path(i, j)
        print(s)
